# Copyright (c) HySoP 2011-2024
#
# This file is part of HySoP software.
# See "https://particle_methods.gricad-pages.univ-grenoble-alpes.fr/hysop-doc/"
# for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
try:
import cPickle as pickle
except:
import pickle
import gzip, portalocker, contextlib, os, errno, uuid, warnings
from hysop import __HYSOP_ENABLE_FILELOCKS__
from hysop.tools.decorators import static_vars
from hysop.tools.warning import HysopCacheWarning
machine_id = None
for path in ("/var/lib/dbus/machine-id", "/var/lib/yum/uuid"):
try:
with open(path) as f:
machine_id = f.read().replace("\n", "")
except:
pass
if machine_id is not None:
break
if machine_id in (None, ""):
machine_id = uuid.getnode()
[docs]
@contextlib.contextmanager
@static_vars(ignored_locks=set())
def lock_file(
filepath,
mode,
compressed=True,
timeout=10,
check_interval=0.1,
ignore_lock_after_timeout=True,
):
"""
Opens a locked file with specified mode, possibly compressed.
"""
_dir = os.path.dirname(filepath)
if not os.path.isdir(_dir):
try:
os.makedirs(_dir)
except OSError as e:
if e.errno != errno.EEXIST:
raise
if not os.path.exists(filepath):
open(filepath, "a").close()
try:
if not __HYSOP_ENABLE_FILELOCKS__ or (
ignore_lock_after_timeout and filepath in lock_file.ignored_locks
):
raise portalocker.exceptions.LockException
with portalocker.Lock(
filename=filepath, timeout=timeout, mode=mode, check_interval=check_interval
) as fl:
if compressed:
with gzip.GzipFile(fileobj=fl, mode=mode) as f:
yield f
else:
yield fl
except portalocker.exceptions.LockException as e:
# Could not obtain the lock in time, so do it the dirty way.
if ignore_lock_after_timeout:
if __HYSOP_ENABLE_FILELOCKS__:
msg = (
f"Could not obtain lock for file '{filepath}' after waiting for {timeout}s, ignoring file lock."
"\nIf this causes a performance issue, consider disabling file locking mechanism by setting "
"environment variable HYSOP_ENABLE_FILELOCKS=0 or by passing --disable-file-locks to your script.\n"
)
warnings.warn(msg, HysopCacheWarning)
lock_file.ignored_locks.add(filepath)
with open(filepath, mode=mode) as fl:
if compressed:
with gzip.GzipFile(fileobj=fl, mode=mode) as f:
yield f
else:
yield fl
else:
msg = f"Could not obtain lock for file '{filepath}' after waiting for {timeout}s, ignoring file lock.\n"
print(f"\nFATAL ERROR: {msg}")
raise e
[docs]
@contextlib.contextmanager
def read_only_lock(filepath, compressed=True, **kwds):
"""Opens a locked read only file, possibly compressed."""
with lock_file(filepath=filepath, mode="rb", compressed=compressed, **kwds) as f:
yield f
[docs]
@contextlib.contextmanager
def write_only_lock(filepath, compressed=True, **kwds):
"""Opens a locked write only file, possibly compressed."""
with lock_file(filepath=filepath, mode="wb", compressed=compressed, **kwds) as f:
yield f
[docs]
def load_cache(filepath, match_type=dict, on_fail={}, **kwds):
"""Load pickled data from filepath atomically."""
data = on_fail
with read_only_lock(filepath, **kwds) as f:
try:
data = pickle.load(f)
if not isinstance(data, match_type):
raise pickle.UnpicklingError
except (OSError, EOFError, pickle.UnpicklingError, AttributeError, TypeError):
data = on_fail
return data
[docs]
def update_cache(filepath, key, data, match_type=dict, on_fail={}, **kwds):
"""
Update cache entry in given file atomically with a (key,data) pair.
Cached data is a pickled dictionnary.
"""
cached_data = load_cache(
filepath=filepath, match_type=match_type, on_fail=on_fail, **kwds
)
cached_data[key] = data
with write_only_lock(filepath=filepath, **kwds) as f:
pickle.dump(cached_data, f)
[docs]
def load_data_from_cache(filepath, key, match_type=dict, on_fail={}, **kwds):
"""Load cached data from a given file atomically with given key."""
data = load_cache(filepath=filepath, match_type=match_type, on_fail=on_fail)
if key in data:
return data[key]
else:
return None
[docs]
def load_attributes_from_cache(filepath, key, instance, attrs, **kwds):
"""
Load cached entries from a given file atomically.
Cached data is assumed to be a dictionnary.
If key is present in pickled data, try to get all
given attributes data by keys given in attrs.
Set instance attributes with those values.
If one attribute is missing or key is not present
in loaded data, set all values to None in instance
and return False.
Return True on success.
"""
success = False
data = load_cache(filepath=filepath, match_type=dict, **kwds)
if key in data:
success = True
data = data[key]
for attr in attrs:
if attr not in data.keys():
for attr in attrs:
setattr(instance, attr, None)
success = False
break
setattr(instance, attr, data[attr])
return success
